package com.cmb.lk50.common.extractor;

import io.netty.buffer.UnpooledByteBufAllocator;
import org.springframework.core.ResolvableType;
import org.springframework.core.codec.ByteArrayDecoder;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.codec.DecoderHttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Author       : 蒋俊钊 80256973
 * Date         : 2019/11/13
 * Copyright    (C) ChinaMerchantsBank
 */
public class BodyParamExtractor extends QueryParamExtractor {

    // body报文提取的扩展点
    protected String extractBody(String body){
        return body;
    }

    @Override
    public String extract(HttpServletRequest request) {
        RequestBodyExtractorWrapper requestWrapper;
        try {
            requestWrapper = new RequestBodyExtractorWrapper(request);
            return extractBody(requestWrapper.getBody());

        } catch (IOException e) {
            throw new RuntimeException("BodyParamExtractor读取body异常");
        }
    }

    /**
     * 封装原来的HttpServletRequest，将其输入流里的数据保存起来
     * 重写getInputStream方法，使其后面读取时都是从保存起来的数据中读的
     */
    public class RequestBodyExtractorWrapper extends HttpServletRequestWrapper {
        private final String body;

        public RequestBodyExtractorWrapper(HttpServletRequest request) throws IOException {
            super(request);

            StringBuilder sb = new StringBuilder();
            try (InputStream inputStream = request.getInputStream()) {
                BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
                String line;
                while ((line = reader.readLine()) != null) {
                    sb.append(line);
                }
            }
            body = sb.toString();
        }

        public String getBody() {
            return body;
        }

        @Override
        public BufferedReader getReader() throws IOException {
            return new BufferedReader(new InputStreamReader(getInputStream()));
        }

        @Override
        public ServletInputStream getInputStream() throws IOException {
            ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(body.getBytes(StandardCharsets.UTF_8));
            return new ServletInputStream() {

                @Override
                public int read() throws IOException {
                    return byteArrayInputStream.read();
                }

                @Override
                public boolean isFinished() {
                    return false;
                }

                @Override
                public boolean isReady() {
                    return false;
                }

                @Override
                public void setReadListener(ReadListener readListener) {
                }
            };
        }
    }

    //////////////////////////////////////////////////////////////////////////


    /**
     * Webflux下的body提取
     * 只在 2.0.6.RELEASE + Finchley.SR2 版本中有效
     */
    @Override
    public String extract(ServerHttpRequest request) {
        DecoderHttpMessageReader<byte[]> reader = new DecoderHttpMessageReader<>(new ByteArrayDecoder());
        ResolvableType resolvableType = ResolvableType.forClass(byte[].class);
        Mono<byte[]> mono = reader.readMono(resolvableType, request, Collections.emptyMap());

        AtomicReference<String> s1 = new AtomicReference<>();
        mono.map(String::new).subscribe(s1::set);

        return extractBody(s1.get());
    }


    public class WrapperServerHttpRequestDecorator extends ServerHttpRequestDecorator {
        private String bodyStr;

        private NettyDataBufferFactory bufferFactory = new NettyDataBufferFactory(new UnpooledByteBufAllocator(false));

        public WrapperServerHttpRequestDecorator(ServerHttpRequest delegate) {
            super(delegate);
        }

        @Override
        public Flux<DataBuffer> getBody() {
            if (bodyStr == null) {
                return super.getBody().map(this::wrapperBody);
            } else {
                return Flux.just(bufferFactory.wrap(bodyStr.getBytes()));
            }
        }

        // 将body中的内容保留下来
        private DataBuffer wrapperBody(DataBuffer buffer) {
            bodyStr = buffer.toString(Charset.defaultCharset());
            DataBufferUtils.release(buffer);
            return bufferFactory.wrap(bodyStr.getBytes());
        }
    }
}
